#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将FrankMocap顶点动画转换为骨骼动画数据
输出标准的骨骼动画pkl文件，可直接导入Blender
"""

import os
import pickle
import numpy as np
from tqdm import tqdm
import argparse

class VertexToBoneAnimationConverter:
    def __init__(self):
        # SMPL/SMPLX骨骼定义（24个主要关节）
        self.joint_names = [
            'Pelvis',           # 0
            'L_Hip',           # 1 
            'R_Hip',           # 2
            'Spine1',          # 3
            'L_Knee',          # 4
            'R_Knee',          # 5
            'Spine2',          # 6
            'L_Ankle',         # 7
            'R_Ankle',         # 8
            'Spine3',          # 9
            'L_Foot',          # 10
            'R_Foot',          # 11
            'Neck',            # 12
            'L_Collar',        # 13
            'R_Collar',        # 14
            'Head',            # 15
            'L_Shoulder',      # 16
            'R_Shoulder',      # 17
            'L_Elbow',         # 18
            'R_Elbow',         # 19
            'L_Wrist',         # 20
            'R_Wrist',         # 21
            'L_Hand',          # 22
            'R_Hand'           # 23
        ]
        
        # 骨骼层级关系
        self.bone_parents = {
            'Pelvis': -1,
            'L_Hip': 0,
            'R_Hip': 0,
            'Spine1': 0,
            'L_Knee': 1,
            'R_Knee': 2,
            'Spine2': 3,
            'L_Ankle': 4,
            'R_Ankle': 5,
            'Spine3': 6,
            'L_Foot': 7,
            'R_Foot': 8,
            'Neck': 9,
            'L_Collar': 9,
            'R_Collar': 9,
            'Head': 12,
            'L_Shoulder': 13,
            'R_Shoulder': 14,
            'L_Elbow': 16,
            'R_Elbow': 17,
            'L_Wrist': 18,
            'R_Wrist': 19,
            'L_Hand': 20,
            'R_Hand': 21
        }
        
        # SMPL关节回归器（简化版本 - 从顶点计算关节位置）
        self.joint_regressor_indices = self.create_joint_regressor()

    def create_joint_regressor(self):
        """创建简化的关节回归器"""
        # 这是SMPL模型中关键顶点的近似索引
        # 实际应用中应该使用官方的joint_regressor矩阵
        regressor = {
            'Pelvis': [3021, 3022, 3023],  # 骨盆区域顶点
            'L_Hip': [1775, 1776, 1777],   # 左髋部
            'R_Hip': [5255, 5256, 5257],   # 右髋部
            'Spine1': [3500, 3501, 3502],  # 脊柱1
            'L_Knee': [1060, 1061, 1062],  # 左膝
            'R_Knee': [4540, 4541, 4542],  # 右膝
            'Spine2': [3800, 3801, 3802],  # 脊柱2
            'L_Ankle': [3365, 3366, 3367], # 左踝
            'R_Ankle': [6845, 6846, 6847], # 右踝
            'Spine3': [4100, 4101, 4102],  # 脊柱3
            'L_Foot': [3392, 3393, 3394],  # 左脚
            'R_Foot': [6872, 6873, 6874],  # 右脚
            'Neck': [412, 413, 414],       # 颈部
            'L_Collar': [1200, 1201, 1202], # 左锁骨
            'R_Collar': [4680, 4681, 4682], # 右锁骨
            'Head': [411, 412, 413],       # 头部
            'L_Shoulder': [1861, 1862, 1863], # 左肩
            'R_Shoulder': [5341, 5342, 5343], # 右肩
            'L_Elbow': [1938, 1939, 1940],    # 左肘
            'R_Elbow': [5418, 5419, 5420],    # 右肘
            'L_Wrist': [2112, 2113, 2114],    # 左腕
            'R_Wrist': [5592, 5593, 5594],    # 右腕
            'L_Hand': [2108, 2109, 2110],     # 左手
            'R_Hand': [5588, 5589, 5590]      # 右手
        }
        
        # 确保索引不超出范围
        max_vertices = 10475  # SMPL-X顶点数
        for joint, indices in regressor.items():
            regressor[joint] = [i % max_vertices for i in indices]
        
        return regressor

    def compute_joint_positions(self, vertices):
        """从顶点计算关节位置"""
        joint_positions = np.zeros((len(self.joint_names), 3))
        
        for i, joint_name in enumerate(self.joint_names):
            if joint_name in self.joint_regressor_indices:
                indices = self.joint_regressor_indices[joint_name]
                # 使用相关顶点的平均位置作为关节位置
                if len(indices) > 0 and max(indices) < len(vertices):
                    joint_positions[i] = np.mean(vertices[indices], axis=0)
                else:
                    # 回退到基于身体比例的估算
                    joint_positions[i] = self.estimate_joint_position(vertices, joint_name)
            else:
                joint_positions[i] = self.estimate_joint_position(vertices, joint_name)
        
        return joint_positions

    def estimate_joint_position(self, vertices, joint_name):
        """基于身体比例估算关节位置"""
        # 计算顶点边界框
        min_coords = np.min(vertices, axis=0)
        max_coords = np.max(vertices, axis=0)
        center = (min_coords + max_coords) / 2
        
        height = max_coords[2] - min_coords[2]
        width = max_coords[0] - min_coords[0]
        depth = max_coords[1] - min_coords[1]
        
        # 基于标准人体比例的关节位置估算
        positions = {
            'Pelvis': [center[0], center[1], min_coords[2] + height * 0.5],
            'L_Hip': [center[0] - width * 0.15, center[1], min_coords[2] + height * 0.48],
            'R_Hip': [center[0] + width * 0.15, center[1], min_coords[2] + height * 0.48],
            'Spine1': [center[0], center[1], min_coords[2] + height * 0.6],
            'L_Knee': [center[0] - width * 0.15, center[1], min_coords[2] + height * 0.25],
            'R_Knee': [center[0] + width * 0.15, center[1], min_coords[2] + height * 0.25],
            'Spine2': [center[0], center[1], min_coords[2] + height * 0.7],
            'L_Ankle': [center[0] - width * 0.15, center[1], min_coords[2] + height * 0.05],
            'R_Ankle': [center[0] + width * 0.15, center[1], min_coords[2] + height * 0.05],
            'Spine3': [center[0], center[1], min_coords[2] + height * 0.8],
            'L_Foot': [center[0] - width * 0.15, center[1] + depth * 0.3, min_coords[2]],
            'R_Foot': [center[0] + width * 0.15, center[1] + depth * 0.3, min_coords[2]],
            'Neck': [center[0], center[1], min_coords[2] + height * 0.9],
            'L_Collar': [center[0] - width * 0.1, center[1], min_coords[2] + height * 0.85],
            'R_Collar': [center[0] + width * 0.1, center[1], min_coords[2] + height * 0.85],
            'Head': [center[0], center[1], min_coords[2] + height * 0.95],
            'L_Shoulder': [center[0] - width * 0.2, center[1], min_coords[2] + height * 0.8],
            'R_Shoulder': [center[0] + width * 0.2, center[1], min_coords[2] + height * 0.8],
            'L_Elbow': [center[0] - width * 0.3, center[1], min_coords[2] + height * 0.65],
            'R_Elbow': [center[0] + width * 0.3, center[1], min_coords[2] + height * 0.65],
            'L_Wrist': [center[0] - width * 0.35, center[1], min_coords[2] + height * 0.5],
            'R_Wrist': [center[0] + width * 0.35, center[1], min_coords[2] + height * 0.5],
            'L_Hand': [center[0] - width * 0.4, center[1], min_coords[2] + height * 0.5],
            'R_Hand': [center[0] + width * 0.4, center[1], min_coords[2] + height * 0.5]
        }
        
        return np.array(positions.get(joint_name, center))

    def compute_bone_rotations(self, joint_positions_current, joint_positions_previous):
        """计算骨骼旋转（基于关节位移）"""
        rotations = np.zeros((len(self.joint_names), 4))  # 四元数 [w, x, y, z]
        
        for i, joint_name in enumerate(self.joint_names):
            if joint_positions_previous is not None:
                # 计算关节移动向量
                movement = joint_positions_current[i] - joint_positions_previous[i]
                movement_magnitude = np.linalg.norm(movement)
                
                if movement_magnitude > 0.001:  # 如果有显著移动
                    # 简化的旋转计算
                    # 实际应用中应该使用更复杂的IK算法
                    rotation_angle = movement_magnitude * 0.1
                    axis = movement / movement_magnitude
                    
                    # 转换为四元数
                    w = np.cos(rotation_angle / 2)
                    sin_half = np.sin(rotation_angle / 2)
                    rotations[i] = [w, axis[0] * sin_half, axis[1] * sin_half, axis[2] * sin_half]
                else:
                    rotations[i] = [1, 0, 0, 0]  # 无旋转
            else:
                rotations[i] = [1, 0, 0, 0]  # 初始帧无旋转
        
        return rotations

    def convert_animation_data(self, input_path, output_path):
        """转换动画数据"""
        print(f"Loading animation data from: {input_path}")
        
        # 加载原始数据
        try:
            with open(input_path, 'rb') as f:
                data = pickle.load(f)
        except Exception as e:
            print(f"Error loading input file: {e}")
            return False
        
        # 提取顶点数据
        if 'vertices' in data:
            vertices_data = data['vertices']  # (frames, vertices, 3)
            num_frames = data.get('num_frames', len(vertices_data))
        else:
            print("No vertex data found in input file")
            return False
        
        print(f"Processing {num_frames} frames...")
        
        # 初始化骨骼动画数据
        bone_animation = {
            'joint_names': self.joint_names,
            'bone_parents': [self.bone_parents.get(name, -1) for name in self.joint_names],
            'num_frames': num_frames,
            'fps': data.get('fps', 30.0),
            'joint_positions': np.zeros((num_frames, len(self.joint_names), 3)),
            'joint_rotations': np.zeros((num_frames, len(self.joint_names), 4)),
            'root_translation': np.zeros((num_frames, 3)),
            'format': 'bone_animation'
        }
        
        previous_joint_positions = None
        initial_root_position = None
        
        # 处理每一帧
        for frame_idx in tqdm(range(num_frames), desc="Converting frames"):
            vertices = vertices_data[frame_idx]
            
            # 计算关节位置
            joint_positions = self.compute_joint_positions(vertices)
            bone_animation['joint_positions'][frame_idx] = joint_positions
            
            # 计算根部平移
            if initial_root_position is None:
                initial_root_position = joint_positions[0].copy()  # Pelvis位置
                bone_animation['root_translation'][frame_idx] = [0, 0, 0]
            else:
                bone_animation['root_translation'][frame_idx] = joint_positions[0] - initial_root_position
            
            # 计算骨骼旋转
            rotations = self.compute_bone_rotations(joint_positions, previous_joint_positions)
            bone_animation['joint_rotations'][frame_idx] = rotations
            
            previous_joint_positions = joint_positions.copy()
        
        # 添加元数据
        bone_animation['original_data'] = {
            'demo_type': data.get('demo_type', 'frank'),
            'smpl_type': data.get('smpl_type', 'smplx'),
            'total_duration': data.get('total_duration', num_frames / 30.0)
        }
        
        # 保存骨骼动画数据
        print(f"Saving bone animation to: {output_path}")
        
        # 创建输出目录
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        
        try:
            with open(output_path, 'wb') as f:
                pickle.dump(bone_animation, f, protocol=2)  # 兼容Blender
        except Exception as e:
            print(f"Error saving output file: {e}")
            return False
        
        # 打印摘要
        print(f"\n{'='*50}")
        print("骨骼动画转换完成!")
        print(f"{'='*50}")
        print(f"输出文件: {output_path}")
        print(f"帧数: {num_frames}")
        print(f"关节数: {len(self.joint_names)}")
        print(f"持续时间: {num_frames/30.0:.2f}秒")
        print(f"文件大小: {os.path.getsize(output_path) / (1024*1024):.2f} MB")
        
        print(f"\n包含的动画数据:")
        print(f"  - joint_positions: {bone_animation['joint_positions'].shape}")
        print(f"  - joint_rotations: {bone_animation['joint_rotations'].shape}")
        print(f"  - root_translation: {bone_animation['root_translation'].shape}")
        
        print(f"\n🎯 现在可以在Blender中使用这个骨骼动画文件!")
        
        return True

def main():
    parser = argparse.ArgumentParser(description='Convert vertex animation to bone animation')
    parser.add_argument('--input', '-i', required=True,
                       help='Input pkl file with vertex animation')
    parser.add_argument('--output', '-o', 
                       default='./output_pkl/bone_animation.pkl',
                       help='Output pkl file with bone animation')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.input):
        print(f"Input file not found: {args.input}")
        return
    
    converter = VertexToBoneAnimationConverter()
    success = converter.convert_animation_data(args.input, args.output)
    
    if success:
        print("\n✅ 转换成功!")
        print("你现在可以将生成的pkl文件导入Blender，它包含标准的骨骼动画数据")
    else:
        print("\n❌ 转换失败!")

if __name__ == "__main__":
    main()